import os
from datetime import datetime

from keras import Sequential, losses
from keras.callbacks import TensorBoard
from keras.layers import Flatten, Dense
from keras import optimizers

from prepare_mnist_data import load_mnist_data, process_mnist_data

if __name__ == "__main__":
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    (train_images, train_labels), (test_images, test_labels) = load_mnist_data()

    train_images = train_images / 255
    test_images = test_images / 255

    model = Sequential(
        [
            Flatten(input_shape=(28, 28)),
            Dense(units=256, activation='relu'),
            Dense(units=128,activation='relu'),
            Dense(units= 32, activation='relu'),
            Dense(units= 10, activation='softmax'),
        ]
    )

    model.compile(optimizer=optimizers.Adam(learning_rate=0.01), loss=losses.sparse_categorical_crossentropy,
                  metrics=['accuracy'])

    model.build(input_shape=[None, 28 * 28])
    model.summary()

    # 定义日志的保存位置
    log_dir = 'logs/fit/' + datetime.now().strftime("%Y%m%d-%H%M%S")
    # 定义 TensorBoard 回调函数
    tensorboard_callback = TensorBoard(log_dir=log_dir, histogram_freq=1)

    model.fit(x=train_images,y=train_labels, epochs=100, callbacks=tensorboard_callback)


